package first_error

import (
	"github.com/panjf2000/ants/v2"
	"sync"
	"sync/atomic"
)

// MyWaitGroup   所有任务执行完成或有一个任务返回error则停止阻塞
type MyWaitGroup struct {
	wg    *sync.WaitGroup
	count *atomic.Int32
	err   error
}

func NewMyWaitGroup() *MyWaitGroup {
	var wg0 sync.WaitGroup
	wg0.Add(1)

	var c atomic.Int32
	return &MyWaitGroup{
		wg:    &wg0,
		count: &c,
	}
}

func (g *MyWaitGroup) Go(f func() error) {
	g.Add(1)
	e := f()
	if e != nil {
		g.err = e
		curCount := g.count.Load()
		g.Add(-curCount)
	} else {
		g.Done()
	}

}

func (g *MyWaitGroup) GoWithAnts(f func() error, pool *ants.Pool) {
	g.Add(1)
	pool.Submit(func() {
		e := f()
		if e != nil {
			g.err = e
			curCount := g.count.Load()
			g.Add(-curCount)
		} else {
			g.Done()
		}
	})

}

func (g *MyWaitGroup) Done() {
	g.Add(-1)
}

func (g *MyWaitGroup) Add(delta int32) {
	curCount := g.count.Add(delta)
	if curCount <= 0 {
		g.wg.Done()
	}
}

// Wait 所有任务执行完成或有一个任务返回error则停止阻塞
func (g *MyWaitGroup) Wait() {
	g.wg.Wait()
}

func (g MyWaitGroup) GetErr() error {
	return g.err
}
